package pcaphelper

import (
	"context"
	"os"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/pcap"
	"github.com/google/gopacket/pcapgo"
	"github.com/pkg/errors"
)

// OnRecvHook 捕获到包的钩子函数
type OnRecvHook func(pkt gopacket.Packet) (isContinue bool)

// ErrNoPktCaptured 没有数据包被捕获
var ErrNoPktCaptured = errors.New("no packet was captured")

// SniffingCtx 嗅探，当收到网络包就调用 onRecv
func (p *PcapHelper) SniffingCtx(ctx context.Context, onRecv OnRecvHook) bool {
	// 强制用户每一次使用的时候先设置BPF
	defer p.SetDireAndBPF(pcap.DirectionIn, "ether dst 0000.0000.0000 and ether src 0000.0000.0000")

	source := gopacket.NewPacketSource(p.handle, p.handle.LinkType())
	source.Lazy = false
	source.DecodeStreamsAsDatagrams = false
	source.NoCopy = true

	in := source.Packets()
	for {
		select {
		case <-ctx.Done():
			return true

		case pkt := <-in:
			// 返回 false 会结束捕获过程
			if isContinue := onRecv(pkt); !isContinue {
				return false
			}
		}
	}
}

// SniffingOnDuration 设置捕获时长
func (p *PcapHelper) SniffingOnDuration(dur time.Duration, onRecv OnRecvHook) bool {
	ctx, _ := context.WithTimeout(context.Background(), dur)
	return p.SniffingCtx(ctx, onRecv)
}

// SniffingOnDurationToFile 设置捕获时长，并将捕获的数据包保存到文件
func (p *PcapHelper) SniffingOnDurationToFile(dur time.Duration, toFile string) (err error) {
	f, err := os.Create(toFile)
	if err != nil {
		return err
	}
	defer f.Close()

	w := pcapgo.NewWriter(f)
	if err := w.WriteFileHeader(uint32(p.handle.SnapLen()), p.handle.LinkType()); err != nil {
		return errors.Wrap(err, "failed to write pcap file header")
	}

	var err1 error
	ctx, _ := context.WithTimeout(context.Background(), dur)
	p.SniffingCtx(ctx, func(pkt gopacket.Packet) bool {
		if err := w.WritePacket(pkt.Metadata().CaptureInfo, pkt.Data()); err != nil {
			err1 = errors.Wrap(err, "an error occurs when writing packet")
			return false
		}
		return true
	})

	return err1
}

// CanCaptureInOne 在指定的时间内能够捕获一个进来的数据包就返回 nil
func (p *PcapHelper) CanCaptureIncoming(dur time.Duration, bpfExpr string) error {
	if err := p.SetDireAndBPF(pcap.DirectionIn, bpfExpr); err != nil {
		return err
	}

	var wasCaptured bool
	p.SniffingOnDuration(dur, func(pkt gopacket.Packet) (isContinue bool) {
		wasCaptured = true
		return false
	})

	if wasCaptured {
		return nil
	}
	return ErrNoPktCaptured
}
